/**
* Copyright 2020 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at

* http://www.apache.org/licenses/LICENSE-2.0

* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.

* File sample_process.cpp
* Description: handle acl resource
*/
#include "yolo/object_detection.h"
#include <iostream>
#include "acl/acl.h"
#include "utils.h"
#include <opencv2/videoio.hpp>

using namespace std;

namespace {
    enum BBoxIndex { TOPLEFTX = 0, TOPLEFTY, BOTTOMRIGHTX, BOTTOMRIGHTY, SCORE, LABEL };
}

YoloObjectDetect::YoloObjectDetect(const char* modelPath,
uint32_t modelWidth,
uint32_t modelHeight
)
:deviceId_(0), context_(nullptr), stream_(nullptr), modelWidth_(modelWidth),
modelHeight_(modelHeight), isInited_(false){
    modelPath_ = modelPath;
}

YoloObjectDetect::~YoloObjectDetect() {
    DestroyResource();
}


Result YoloObjectDetect::InitModel(const char* omModelPath) {
    Result ret = model_.LoadModelFromFileWithMem(omModelPath);
    if (ret != SUCCESS) {
        ERROR_LOG("execute LoadModelFromFileWithMem failed");
        return FAILED;
    }

    ret = model_.CreateDesc();
    if (ret != SUCCESS) {
        ERROR_LOG("execute CreateDesc failed");
        return FAILED;
    }

    ret = model_.CreateOutput();
    if (ret != SUCCESS) {
        ERROR_LOG("execute CreateOutput failed");
        return FAILED;
    }

    return SUCCESS;
}

Result YoloObjectDetect::Init(int32_t& deviceId, aclrtContext& context, aclrtStream& stream, aclrtRunMode& runMode) {
    if (isInited_) {
        INFO_LOG("yolo instance is initied already!");
        return SUCCESS;
    }


    deviceId_ = deviceId;
    context_ = context;
    stream_ = stream;

    runMode_ = runMode;

    aclError ret = InitModel(modelPath_);
    if (ret != SUCCESS) {
        ERROR_LOG("Init model failed");
        return FAILED;
    }

    ret = dvpp_.InitResource(stream_);
    if (ret != SUCCESS) {
        ERROR_LOG("Init dvpp failed");
        return FAILED;
    }

    ret = CreateYolov5InputdDataset();
    if (ret != SUCCESS) {
        ERROR_LOG("Create image info buf failed");
        return FAILED;
    }

    isInited_ = true;
    return SUCCESS;
}

Result YoloObjectDetect::CreateYolov3InputdDataset()
{
    //Request image data memory for input model
    imageDataSize_ = RGBU8_IMAGE_SIZE(modelWidth_, modelHeight_);
    aclError aclRet = aclrtMalloc(&imageDataBuf_, imageDataSize_, ACL_MEM_MALLOC_HUGE_FIRST);
    if (aclRet != ACL_ERROR_NONE) {
        ERROR_LOG("malloc device data buffer failed, aclRet is %d", aclRet);
        return FAILED;
    }

    //Use the applied memory to create the model and input dataset. After creation, only update the memory data for each frame of inference, instead of creating the input dataset every time
    Result ret = model_.CreateInput(imageDataBuf_, imageDataSize_);
    if (ret != SUCCESS) {
        ERROR_LOG("Create model input dataset failed");
        return FAILED;
    }

    return SUCCESS;
}

Result YoloObjectDetect::CreateYolov5InputdDataset()
{
    //Request image data memory for input model
    imageDataSize_ = RGBU8_IMAGE_SIZE(modelWidth_, modelHeight_);
    aclError aclRet = aclrtMalloc(&imageDataBuf_, imageDataSize_, ACL_MEM_MALLOC_HUGE_FIRST);
    if (aclRet != ACL_ERROR_NONE) {
        ERROR_LOG("malloc device data buffer failed, aclRet is %d", aclRet);
        return FAILED;
    }

    //Use the applied memory to create the model and input dataset. After creation, only update the memory data for each frame of inference, instead of creating the input dataset every time
    const float imageInfo[4] = {(float)modelWidth_, (float)modelHeight_,
    (float)modelWidth_, (float)modelHeight_};
    imageInfoSize_ = sizeof(imageInfo);
    if (runMode_ == ACL_HOST)
        imageInfoBuf_ = Utils::CopyDataHostToDevice((void *)imageInfo, imageInfoSize_);
    else
        imageInfoBuf_ = Utils::CopyDataDeviceToDevice((void *)imageInfo, imageInfoSize_);
    if (imageInfoBuf_ == nullptr) {
        ERROR_LOG("Copy image info to device failed");
        return FAILED;
    }
    Result ret = model_.CreateInput(imageDataBuf_, imageDataSize_, imageInfoBuf_, imageInfoSize_);
    if (ret != SUCCESS) {
        ERROR_LOG("Create model input dataset failed");
        return FAILED;
    }

    return SUCCESS;
}

// ImageData为YUV420spNV12格式的图像，Mat为bgr格式的空图像，vector为检测结果
Result YoloObjectDetect::Process(ImageData& image, cv::Mat& bgrimg, vector<BBox>& dectections) {
    // //预处理图片:读取图片,讲图片缩放到模型输入要求的尺寸
    
    struct timeval begin;
    struct timeval end;
    gettimeofday(&begin, NULL);
    Result ret = Preprocess(image, bgrimg);
    gettimeofday(&end, NULL);
    if (ret != SUCCESS) {
        ERROR_LOG("Preprocess failed");      
        return FAILED;
    }
    // float time_cost = 1000 * (end.tv_sec - begin.tv_sec) + (end.tv_usec - begin.tv_usec) / 1000.000;
    // INFO_LOG("Preprocess time: %fms", time_cost);


    //将预处理的图片送入模型推理,并获取推理结果
    aclmdlDataset* inferenceOutput = nullptr;
    gettimeofday(&begin, NULL);
    ret = Inference(inferenceOutput);
    gettimeofday(&end, NULL);
    if ((ret != SUCCESS) || (inferenceOutput == nullptr)) {
        ERROR_LOG("Inference model inference output data failed");
        return FAILED;
    }
    // time_cost = 1000 * (end.tv_sec - begin.tv_sec) + (end.tv_usec - begin.tv_usec) / 1000.000;
    // INFO_LOG("Inference time: %fms", time_cost);

    //解析推理输出,并将推理得到的物体类别和位置标记到图片上
    gettimeofday(&begin, NULL);
    ret = Postprocess(inferenceOutput, bgrimg, dectections);
    gettimeofday(&end, NULL);
    if (ret != SUCCESS) {
        ERROR_LOG("Process model inference output data failed");
        return FAILED;
    }
    // time_cost = 1000 * (end.tv_sec - begin.tv_sec) + (end.tv_usec - begin.tv_usec) / 1000.000;
    // INFO_LOG("Postprocess time: %fms", time_cost);
    // INFO_LOG("Target Size: %d", dectections.size());
    
    return SUCCESS;
}

Result YoloObjectDetect::Preprocess(ImageData& srcImage, cv::Mat& bgrimg) {
    // INFO_LOG("start Preprocess...");
    Result ret;

    // 将dvpp解码的yuv转为rgb便于后续处理
    cv::Mat yuvimg(srcImage.height * 3 / 2, srcImage.width, CV_8UC1, srcImage.data.get());
    cv::cvtColor(yuvimg, bgrimg, cv::COLOR_YUV2BGR_NV12);

    // 用dvpp做缩放，然后再用opencv做letterbox，图片越大，用这种方式的速度优势就越明显，但是dvpp是会带了一定的精度下降的
    ImageData resizedImageData;
    dvpp_.ResizeWithLetterBox(resizedImageData, srcImage, cv::Size(modelWidth_, modelHeight_));  // 用dvpp缩放
    cv::Mat resizedImage(resizedImageData.height * 3 / 2, resizedImageData.width, CV_8UC1, resizedImageData.data.get());
    cv::cvtColor(resizedImage, resizedImage, cv::COLOR_YUV2BGR_NV12);
    cv::Mat letterBoxImage = Utils::letterbox_img(resizedImage, cv::Size(modelWidth_, modelHeight_)); // 将图像等比例缩放为模型输入大小


    // 直接用opencv缩放和做letterbox
    // cv::Mat letterBoxImage = Utils::letterbox_img(bgrimg, cv::Size(modelWidth_, modelHeight_)); // 将图像等比例缩放为模型输入大小

    //Copy the data into the cache of the input dataset
    aclrtMemcpyKind policy = (runMode_ == ACL_HOST)?
                             ACL_MEMCPY_HOST_TO_DEVICE:ACL_MEMCPY_DEVICE_TO_DEVICE;
    ret = (Result)aclrtMemcpy(imageDataBuf_, imageDataSize_,
                               letterBoxImage.data, imageDataSize_, policy);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("Copy resized image data to device failed.");
        return FAILED;
    }

    const float imageInfo[4] = {(float)modelHeight_, (float)modelWidth_,
    (float)srcImage.height, (float)srcImage.width};
    ret = (Result)aclrtMemcpy(imageInfoBuf_, imageInfoSize_,
                               (void *)imageInfo, imageInfoSize_, policy);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("Copy imageinfo to device failed.");
        return FAILED;
    }


    return SUCCESS;
}

Result YoloObjectDetect::Inference(aclmdlDataset*& inferenceOutput) {
    // INFO_LOG("start Inference...");

    Result ret = model_.Execute();
    if (ret != SUCCESS) {
        ERROR_LOG("Execute model inference failed");
        return FAILED;
    }

    inferenceOutput = model_.GetModelOutputData();

    // 调试用
    // string modelName = "yolov5s_sim_t";
    // model_.OutputModelResult(modelName);
    return SUCCESS;
}

Result YoloObjectDetect::Postprocess(aclmdlDataset* modelOutput, cv::Mat& bgrimg, vector<BBox>& result)
{
    // // INFO_LOG("start Postprocess...");

    // Result ret;

    // // Get data
    // uint32_t dataSize = 0;
    // float* outputData = (float*)GetInferenceOutputItem(dataSize, modelOutput,
    //                                                    0);
    // if (outputData == nullptr) return FAILED;

    // PostProcess postProcess(outputData, modelWidth_, modelHeight_);

    // ret = postProcess.Process(bgrimg, result);
    // if (ret != SUCCESS) {
    //     ERROR_LOG("postProcess failed");
    //     return FAILED;
    // }


    // //If it is the host side, the data is copied from the device and the memory used by the copy is freed
    // if (runMode_ == ACL_HOST) {
    //     delete[]((uint8_t*)outputData);
    // }

    // return SUCCESS;


    // Get box information data
    uint32_t dataSize = 0;
    float* detectData = (float*)GetInferenceOutputItem(dataSize, modelOutput,
                                                       0);
    if (detectData == nullptr) return FAILED;
    //Gets the number of boxes
    uint32_t* boxNum = (uint32_t*)GetInferenceOutputItem(dataSize, modelOutput,
                                                         1);
    if (boxNum == nullptr) return FAILED;

    //Number of boxes The first data is valid
    uint32_t totalBox = boxNum[0];

    for (uint32_t i = 0; i < totalBox; i++) {
        BBox box;
        //get the confidence of the detected object. Anything less than 0.8 is considered invalid
        uint32_t score = uint32_t(detectData[totalBox * SCORE + i] * 100);

        if (score < 50) continue;

        //get the frame coordinates and converts them to the coordinates on the original frame
        box.rect.x = detectData[totalBox * TOPLEFTX + i] ;
        box.rect.y = detectData[totalBox * TOPLEFTY + i] ;
        box.rect.width = detectData[totalBox * BOTTOMRIGHTX + i] - box.rect.x;
        box.rect.height = detectData[totalBox * BOTTOMRIGHTY + i] - box.rect.y;
        box.cls_id = (uint32_t)detectData[totalBox * LABEL + i];
        box.score = score;
        result.emplace_back(box);
    }

    //If it is the host side, the data is copied from the device and the memory used by the copy is freed
    if (runMode_ == ACL_HOST) {
        delete[]((uint8_t*)detectData);
        delete[]((uint8_t*)boxNum);
    }

    return SUCCESS;
}

void YoloObjectDetect::DrawDetections(cv::Mat& ori_img, vector<BBox>& detections) {
    PostProcess::DrawDetections(ori_img, detections);
}

void* YoloObjectDetect::GetInferenceOutputItem(uint32_t& itemDataSize,
aclmdlDataset* inferenceOutput,
uint32_t idx) {
    aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(inferenceOutput, idx);
    if (dataBuffer == nullptr) {
        ERROR_LOG("Get the %dth dataset buffer from model "
        "inference output failed", idx);
        return nullptr;
    }

    void* dataBufferDev = aclGetDataBufferAddr(dataBuffer);
    if (dataBufferDev == nullptr) {
        ERROR_LOG("Get the %dth dataset buffer address "
        "from model inference output failed", idx);
        return nullptr;
    }

    size_t bufferSize = aclGetDataBufferSizeV2(dataBuffer);
    if (bufferSize == 0) {
        ERROR_LOG("The %dth dataset buffer size of "
        "model inference output is 0", idx);
        return nullptr;
    }

    void* data = nullptr;
    if (runMode_ == ACL_HOST) {
        data = Utils::CopyDataDeviceToLocal(dataBufferDev, bufferSize);
        if (data == nullptr) {
            ERROR_LOG("Copy inference output to host failed");
            return nullptr;
        }
    }
    else {
        data = dataBufferDev;
    }

    itemDataSize = bufferSize;
    return data;
}

void YoloObjectDetect::DestroyResource()
{
    model_.Unload();
    model_.DestroyDesc();
    model_.DestroyInput();
    model_.DestroyOutput();
    dvpp_.DestroyResource();

}